from xml.etree import cElementTree
import glob
import shutil


class DataAdjustor:
    @staticmethod
    def combine_class(xml_path, class_names, save_path, combine_name):
        """

        :param combine_name:
        :param xml_path:list:
        :param class_names:list:最后一个类名即是最终修改成的类名
        :param save_path:string:保存路径
        """
        xml_list = glob.glob(xml_path + '\\' + '*.xml')
        for path in xml_list:
            save_path_ = save_path + '\\' + path.split('\\')[-1]
            img_save_path = save_path_ + '.jpg'
            root_element = cElementTree.parse(path)
            root = root_element.getroot()
            for obj in root.iter('object'):
                if obj.find('name').text in class_names:
                    obj.find('name').text = combine_name
            img_name = path.split('.')[0] + '.jpg'
            shutil.copy(img_name, img_save_path)
            root_element.write(save_path_)

    @staticmethod
    def alter_difficult_data(xml_path, difficult_class_names, save_name):
        """

        :param xml_path:
        :param difficult_class_names:
        :param save_name:
        :return:
        """
        xml_list = glob.glob(xml_path + '\\' + '*.xml')
        for path in xml_list:
            save_path = save_name + '\\' + path.split('\\')[-1]
            img_save_path = save_name + '\\' + path.split('\\')[-1].split('.')[0] + '.jpg'
            root_element = cElementTree.parse(path)
            root = root_element.getroot()
            for obj in root.iter('object'):
                if obj.find('name').text in difficult_class_names:
                    obj.find('difficult').text = '1'
                    print(obj.find('difficult').text)
            img_name = path.split('.')[0] + '.jpg'
            shutil.copy(img_name, img_save_path)
            root_element.write(save_path)


def main():
    path = r'D:\test'
    da = DataAdjustor()
    da.combine_class(path, ['gerber23', 'gerber21', 'gerber25', 'gerber28'], r'D:\test_1', 'gerber')


if __name__ == '__main__':
    main()
